package zhengqc.graph.algorithm.mst;


import zhengqc.graph.algorithm.ConnectedComponent;
import zhengqc.graph.data_structure.WeightedEdge;
import zhengqc.graph.data_structure.WeightedGraph;
import zhengqc.queue.PriorityQueue;

import java.util.ArrayList;

/* 暂时仅支持无向有权图 */
public class Prim {
    private WeightedGraph graph;
    private ArrayList<WeightedEdge> mst;
    private int minimalCost;

    public Prim(WeightedGraph graph) {
        this.graph = graph;
        this.minimalCost = 0;
        this.mst = new ArrayList<>();
        ConnectedComponent wcc = new ConnectedComponent(graph);
        if (wcc.getCount() > 1) {
            return;
        }
        boolean[] visited = new boolean[graph.getV()];
        visited[0] = true;
        PriorityQueue<WeightedEdge> pq = new PriorityQueue<>();
        for (int w: graph.getAdjacency(0)) {
            pq.enqueue(new WeightedEdge(0, w, graph.getWeighted(0, w)));
        }
        while (!pq.isEmpty()) {
            WeightedEdge minEdge = pq.dequeue();
            // 若是 minEdge 的两个顶点都属于已划分的部分, 不会执行后续逻辑
            if (visited[minEdge.getV()] && visited[minEdge.getW()]) {
                continue;
            }
            mst.add(minEdge);
            minimalCost += minEdge.getWeight();
            int newV = visited[minEdge.getV()] ? minEdge.getW() : minEdge.getV();
            visited[newV] = true;
            for (int w: graph.getAdjacency(newV)) {
                if (!visited[w]) {
                    pq.enqueue(new WeightedEdge(newV, w, graph.getWeighted(newV, w)));
                }
            }
        }
    }
    public ArrayList<WeightedEdge> getMST() {
        return this.mst;
    }

    public int getMinimalCost() {
        return minimalCost;
    }
}
